package com.ho1ho.audioexample.utils.others;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;

/**
 * Created by Michael Leo <y@ho1ho.com><br>
 * Date: 2020/02/03 17:05
 * <p>
 * Code is far away from bug with the animal protecting
 * </p>
 * <pre>
 * ----------Dragon be here!----------/
 *    ┏┓   ┏┓
 *   ┏┛┻━━━┛┻┓
 *   ┃       ┃
 *   ┃   ━   ┃
 *   ┃ ┳┛ ┗┳ ┃
 *   ┃       ┃
 *   ┃   ┻   ┃
 *   ┃       ┃
 *   ┗━┓   ┏━┛
 *     ┃   ┃神兽保佑
 *     ┃   ┃代码无BUG！
 *     ┃   ┗━━━┓
 *     ┃       ┣┓
 *     ┃       ┏┛
 *     ┗┓┓┏━┳┓┏┛
 *      ┃┫┫ ┃┫┫
 *      ┗┻┛ ┗┻┛
 * ━━━━━━神兽出没━━━━━━
 * </pre>
 */
public class PcmToWavUtil {
    /**
     * Size of buffer used for transfer, by default
     */
    private static final int TRANSFER_BUFFER_SIZE = 10 * 1024;

    /**
     * @param pcmData      Raw PCM data
     * @param numChannels  Channel count. mono = 1, stereo = 2
     * @param sampleRate   Sample rate
     * @param bitPerSample Bits per sample. Example: 8bits, 16bits
     * @return Wave data
     */
    public static byte[] pcmToWav(byte[] pcmData, int numChannels, int sampleRate, int bitPerSample) {
        byte[] wavData = new byte[pcmData.length + 44];
        byte[] header = wavHeader(pcmData.length, numChannels, sampleRate, bitPerSample);
        System.arraycopy(header, 0, wavData, 0, header.length);
        System.arraycopy(pcmData, 0, wavData, header.length, pcmData.length);
        return wavData;
    }

    // ====================================================

    /**
     * @param pcmLen       The length of PCM
     * @param numChannels  Channel count. mono = 1, stereo = 2
     * @param sampleRate   Sample rate
     * @param bitPerSample Bits per sample. Example: 8bits, 16bits
     * @return Wave header
     */
    public static byte[] wavHeader(int pcmLen, int numChannels, int sampleRate, int bitPerSample) {
        byte[] header = new byte[44];
        // ChunkID, RIFF, 4 bytes respectively.
        header[0] = 'R';
        header[1] = 'I';
        header[2] = 'F';
        header[3] = 'F';
        // ChunkSize, pcmLen + 36, 4 bytes respectively.
        long chunkSize = pcmLen + 36;
        header[4] = (byte) (chunkSize & 0xff);
        header[5] = (byte) ((chunkSize >> 8) & 0xff);
        header[6] = (byte) ((chunkSize >> 16) & 0xff);
        header[7] = (byte) ((chunkSize >> 24) & 0xff);
        // Format, WAVE, 4 bytes respectively
        header[8] = 'W';
        header[9] = 'A';
        header[10] = 'V';
        header[11] = 'E';
        // Subchunk1 ID, 'fmt ', 4 bytes
        header[12] = 'f';
        header[13] = 'm';
        header[14] = 't';
        header[15] = ' ';
        // Subchunk1 Size, 16, 4 bytes
        header[16] = 16;
        header[17] = 0;
        header[18] = 0;
        header[19] = 0;
        // AudioFormat, pcm = 1, 2bytes
        header[20] = 1;
        header[21] = 0;
        // NumChannels, mono = 1, stereo = 2, 2 bytes
        header[22] = (byte) numChannels;
        header[23] = 0;
        // SampleRate, 4 bytes
        header[24] = (byte) (sampleRate & 0xff);
        header[25] = (byte) ((sampleRate >> 8) & 0xff);
        header[26] = (byte) ((sampleRate >> 16) & 0xff);
        header[27] = (byte) ((sampleRate >> 24) & 0xff);
        // ByteRate = SampleRate * NumChannels * BitsPerSample / 8, 4 bytes
        long byteRate = sampleRate * numChannels * bitPerSample / 8;
        header[28] = (byte) (byteRate & 0xff);
        header[29] = (byte) ((byteRate >> 8) & 0xff);
        header[30] = (byte) ((byteRate >> 16) & 0xff);
        header[31] = (byte) ((byteRate >> 24) & 0xff);
        // BlockAlign = NumChannels * BitsPerSample / 8, 2 bytes
        header[32] = (byte) (numChannels * bitPerSample / 8);
        header[33] = 0;
        // BitsPerSample, 2 bytes
        header[34] = (byte) bitPerSample;
        header[35] = 0;
        // Subhunk2ID, data, 4 bytes
        header[36] = 'd';
        header[37] = 'a';
        header[38] = 't';
        header[39] = 'a';
        // Subchunk2Size, 4 bytes
        header[40] = (byte) (pcmLen & 0xff);
        header[41] = (byte) ((pcmLen >> 8) & 0xff);
        header[42] = (byte) ((pcmLen >> 16) & 0xff);
        header[43] = (byte) ((pcmLen >> 24) & 0xff);

        return header;
    }

    /**
     * @param input         raw PCM data
     *                      limit of file size for wave file: < 2^(2*4) - 36 bytes (~4GB)
     * @param output        file to encode to in wav format
     * @param channelCount  number of channels: 1 for mono, 2 for stereo, etc.
     * @param sampleRate    sample rate of PCM audio
     * @param bitsPerSample bits per sample, i.e. 16 for PCM16
     * @throws IOException in event of an error between input/output files
     * @see <a href="http://soundfile.sapp.org/doc/WaveFormat/">soundfile.sapp.org/doc/WaveFormat</a>
     */
    static public void pcmToWav(File input, File output, int channelCount, int sampleRate, int bitsPerSample) throws IOException {
        final int inputSize = (int) input.length();

        try (OutputStream encoded = new FileOutputStream(output)) {
            // WAVE RIFF header
            writeToOutput(encoded, "RIFF"); // chunk id
            writeToOutput(encoded, 36 + inputSize); // chunk size
            writeToOutput(encoded, "WAVE"); // format

            // SUB CHUNK 1 (FORMAT)
            writeToOutput(encoded, "fmt "); // subchunk 1 id
            writeToOutput(encoded, 16); // subchunk 1 size
            writeToOutput(encoded, (short) 1); // audio format (1 = PCM)
            writeToOutput(encoded, (short) channelCount); // number of channelCount
            writeToOutput(encoded, sampleRate); // sample rate
            writeToOutput(encoded, sampleRate * channelCount * bitsPerSample / 8); // byte rate
            writeToOutput(encoded, (short) (channelCount * bitsPerSample / 8)); // block align
            writeToOutput(encoded, (short) bitsPerSample); // bits per sample

            // SUB CHUNK 2 (AUDIO DATA)
            writeToOutput(encoded, "data"); // subchunk 2 id
            writeToOutput(encoded, inputSize); // subchunk 2 size
            copy(new FileInputStream(input), encoded);
        }
    }

    /**
     * Writes string in big endian form to an output stream
     *
     * @param output stream
     * @param data   string
     * @throws IOException
     */
    public static void writeToOutput(OutputStream output, String data) throws IOException {
        for (int i = 0; i < data.length(); i++)
            output.write(data.charAt(i));
    }

    public static void writeToOutput(OutputStream output, int data) throws IOException {
        output.write(data & 0xff);
        output.write(data >> 8 & 0xff);
        output.write(data >> 16 & 0xff);
        output.write(data >> 24 & 0xff);
    }

    public static void writeToOutput(OutputStream output, short data) throws IOException {
        output.write(data & 0xff);
        output.write(data >> 8 & 0xff);
    }

    public static long copy(InputStream source, OutputStream output)
            throws IOException {
        return copy(source, output, TRANSFER_BUFFER_SIZE);
    }

    public static long copy(InputStream source, OutputStream output, int bufferSize) throws IOException {
        long read = 0L;
        byte[] buffer = new byte[bufferSize];
        for (int n; (n = source.read(buffer)) != -1; read += n) {
            output.write(buffer, 0, n);
        }
        return read;
    }
}
